from langchain_core.exceptions import OutputParserException
from langchain_core.output_parsers import BaseOutputParser
from langchain_community.embeddings.huggingface import HuggingFaceBgeEmbeddings
from langchain_community.document_loaders.json_loader import JSONLoader
# from langchain_community.vectorstores.docarray.in_memory import DocArrayInMemorySearch
from langchain_community.vectorstores.faiss import FAISS

import re
from typing import Literal, Any

CODE_INSTRUCTIONS = {
    "code_trans": """
输出应当被格式化为一个Markdown形式的代码块
一个输出示例如下：
```
(code content, such as rust)
```
""",
    "code_think": """
请你仔细分析源代码的各条语句，一步一步思考应该如何完成代码转换工作，只进行分析，不写代码，各条分析结果之间使用“-------”分割开
一个示例输出如下：
对代码转换任务的思考1：
-------
对代码转换任务的思考2：
-------
对代码转换任务的思考3：
""",
}

class CodeBlockOutputParser(BaseOutputParser[str]):
	# 自定义的一个字段，用于指定返回代码块的类型
    code_type: str = None

	# 在这段初始化方法中，指定了code_type只能为限定类型的代码
    def __init__(self, code_type: Literal["c", "c++", "rust"] = "rust", **kwargs: Any):
        super().__init__(**kwargs)
        self.code_type = code_type

    def parse(self, text: str) -> str:
        """该方法用于解析模型的输出"""
        # 使用非贪婪匹配 .*? 来捕获三个反引号之间的任何内容
        pattern = r'```(.*?)```'

        # 执行搜索
        match = re.search(pattern, text, re.DOTALL)

        # 判断并输出结果
        if match:
            code = match.group(1)
            code = code.strip(self.code_type).strip()
            return code
        else:
            raise OutputParserException("The response has no code block.", llm_output=text)

    @property
    def _type(self) -> str:
        """返回该解析器的类型 这里返回的是自定义代码块解析器"""
        return "MarkdownCodeBlock"

def create_json_retriver(json_file):
    bge_model = load_embedding_model()
    loader = JSONLoader(file_path=json_file, jq_schema=".[]", text_content=False)
    docs = loader.load()
    db = FAISS.from_documents(
        docs,
        bge_model
    )
    return db.as_retriever(k=1)

def load_embedding_model():

    model_name = "BAAI/bge-m3"
    model_kwargs = {"device": "cpu"}
    encode_kwargs = {"normalize_embeddings": True}
    bge_model = HuggingFaceBgeEmbeddings(
        model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
    )
    return bge_model